#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import numpy as np
import face_recognition
from flask import Flask, jsonify, request, redirect, render_template

photo_path = ".\\photo\\"       # 照片路径
known_face_names = []           # 已知面容姓名
known_face_encodings = []       # 已知面容编码
compare_faces_tolerance = 0.39  # 阈值，阈值太低容易造成无法成功识别人脸，太高容易造成人脸识别混淆，默认0.6，实测0.39会更准确一些

# 加载照片资源
def loadFiles(root_path, tab_flag):
    fileList = os.listdir(root_path)
    for file in fileList:
        item_path = os.path.join(root_path, file)
        if os.path.isdir(item_path):
            print(tab_flag + file + " ↓")
            loadFiles(item_path, tab_flag + "····")
        else:
            print(tab_flag + file)
            image = face_recognition.load_image_file(item_path)
            face_encoding = face_recognition.face_encodings(image)        
            if (len(face_encoding) == 0):
                print("< 未检测到人脸，加载失败:( >")
                continue
            name = os.path.splitext(file)[0]
            known_face_encodings.append(face_encoding[0])
            known_face_names.append(name)


print("开始加载照片资源...")
loadFiles(photo_path, "|")
print("照片资源加载完毕，共{}张照片！".format(len(known_face_names)))


app = Flask(__name__)

# 允许上传文件的扩展名
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# 人脸识别测试
@app.route('/', methods=['GET', 'POST'])
def upload_image():
    # 检查是否上传了有效的图像文件
    if request.method == 'POST':
        if 'file' not in request.files:
            return redirect(request.url)

        file = request.files['file']

        if file.filename == '':
            return redirect(request.url)

        if file and allowed_file(file.filename):
            # 图像有效，检测人脸并返回结果
            return detect_faces_in_image(file)

    # 如果没有上传有效的图像，显示文件上传表单
    return render_template("index.html")


# 检测人脸并返回结果
def detect_faces_in_image(file_stream):
    # 加载上传的图像文件
    img = face_recognition.load_image_file(file_stream)
    # 获取上传图像中的人脸编码
    unknown_face_encodings = face_recognition.face_encodings(img)

    face_found = False
    recognition_result = False

    if len(unknown_face_encodings) > 0:
        face_found = True
        # 匹配
        match_results = face_recognition.compare_faces(known_face_encodings, unknown_face_encodings[0], tolerance=compare_faces_tolerance)
        if True in match_results:
            match_index = match_results.index(True)
            recognition_result = known_face_names[match_index]

    # 返回结果
    result = {
        "face_found": face_found,
        "recognition_result": recognition_result
    }
    print(result)
    return jsonify(result)


if __name__ == "__main__":
    app.config['JSON_AS_ASCII'] = False
    app.run(host='0.0.0.0', port=5001, debug=True, use_reloader=False)
